import torch
import subprocess
import json
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from UNetResNet50 import UNetResNet50
from data_loader import Mdataset,transform
def get_gpu_memory_map():
    """Get the GPU memory map through nvidia-smi"""
    result = subprocess.check_output(
        [
            'nvidia-smi', '--query-gpu=memory.free',
            '--format=csv,nounits,noheader'
        ], encoding='utf-8')
    # Convert lines into a dictionary
    gpu_memory = [int(x) for x in result.strip().split('\n')]
    return dict(enumerate(gpu_memory))

def find_best_gpu():
    """Find the GPU with the most free memory."""
    gpu_memory_map = get_gpu_memory_map()
    best_gpu_id = max(gpu_memory_map, key=gpu_memory_map.get)
    print(f"GPU {best_gpu_id} has the most free memory: {gpu_memory_map[best_gpu_id]} MB")
    torch.cuda.set_device(best_gpu_id)
    return best_gpu_id
if __name__ == "__main__":
    find_best_gpu()
    root_dir = '/home/jinyue/work/pcy241003/data/edeme/'
    n_classes = 1  # Binary segmentation
    batch_size = 1 # Adjust based on your GPU memory
    learning_rate = 1e-5
    epochs = 100

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Model, dataset, and data loader
    model = UNetResNet50(n_classes=n_classes).to(device)

    dataset = Mdataset(root_dir=root_dir, transform=transform)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

    # Optimizer and loss function
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    criterion = nn.BCEWithLogitsLoss()

    # Training loop
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        for i, (raw_img, mask_img) in enumerate(train_loader):
            raw_img = raw_img.to(device)
            mask_img = mask_img.to(device).float().unsqueeze(1)

            # Forward pass
            outputs = model(raw_img)
            loss = criterion(outputs, mask_img)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}')

    print('Training complete.')

    # Save the model
    torch.save(model.state_dict(), 'unet_resnet50_model.pth')
    print('Model saved successfully.')
